use super::agents::{AgentManager, ChatMessage};
use super::tools::ToolManager;
use futures::future::Either;
use futures::stream::{Stream, StreamExt};
use rig::client::{CompletionClient, ProviderClient};
use rig::completion::{Chat, Prompt};
use rig::message::UserContent;
use rig::one_or_many::OneOrMany;
use rig::providers::{anthropic, deepseek, gemini, groq, openai, openrouter};
use std::pin::Pin;

macro_rules! with_client {
    ($provider:expr, $base_url:expr, |$client:ident| $block:expr) => {
        match $provider.as_str() {
            "openai" => {
                let client = openai::Client::from_env();
                let $client = client;
                $block
            }
            "anthropic" => {
                let $client = anthropic::Client::from_env();
                $block
            }
            "deepseek" => {
                let $client = deepseek::Client::from_env();
                $block
            }
            "gemini" => {
                let $client = gemini::Client::from_env();
                $block
            }
            "groq" => {
                let $client = groq::Client::from_env();
                $block
            }
            "openrouter" => {
                let $client = openrouter::Client::from_env();
                $block
            }
            _ => Err(anyhow::anyhow!("Unsupported provider: {}", $provider)),
        }
    };
}

pub async fn prompt(
    agent_manager: &AgentManager,
    tool_manager: &ToolManager,
    agent_id: i64,
    prompt: &str,
) -> anyhow::Result<String> {
    // Get the agent to determine its category
    let agent = agent_manager
        .get_agent(agent_id)
        .await
        .map_err(|e| anyhow::anyhow!("Failed to get agent: {}", e))?
        .ok_or_else(|| anyhow::anyhow!("Agent with id {} not found", agent_id))?;

    // Handle regular agents
    with_client!(
        &agent.base_config.provider,
        agent.base_config.base_url.as_deref(),
        |client| {
            let agent_builder = client
                .agent(&agent.base_config.model)
                .preamble(&agent.base_config.preamble);

            // Add tools if available (including MCP tools)
            if let Some(tool_names) = &agent.base_config.tools {
                if !tool_names.is_empty() {
                    let tools = tool_manager
                        .get_agent_all_tools(agent_id)
                        .map_err(|e| anyhow::anyhow!("Failed to get tools: {}", e))?;
                    if !tools.is_empty() {
                        // Note: In a real implementation, you would need to convert ToolDefinition
                        // to rig's tool format. This is a placeholder for the actual integration.
                        // agent_builder = agent_builder.tools(tools);
                    }
                }
            }

            let agent = agent_builder.build();

            let response = agent
                .prompt(prompt)
                .await
                .map_err(|e| anyhow::anyhow!("Failed to prompt agent: {}", e))?;

            Ok(response)
        }
    )
}

pub async fn chat(
    agent_manager: &AgentManager,
    tool_manager: &ToolManager,
    agent_id: i64,
    prompt: &str,
    history: Vec<ChatMessage>,
) -> anyhow::Result<String> {
    // Get the agent to determine its category
    let agent = agent_manager
        .get_agent(agent_id)
        .await
        .map_err(|e| anyhow::anyhow!("Failed to get agent: {}", e))?
        .ok_or_else(|| anyhow::anyhow!("Agent with id {} not found", agent_id))?;

    // Handle regular agents
    with_client!(
        &agent.base_config.provider,
        agent.base_config.base_url.as_deref(),
        |client| {
            let agent_builder = client
                .agent(&agent.base_config.model)
                .preamble(&agent.base_config.preamble);

            // Add tools if available (including MCP tools)
            if let Some(tool_names) = &agent.base_config.tools {
                if !tool_names.is_empty() {
                    let tools = tool_manager
                        .get_agent_all_tools(agent_id)
                        .map_err(|e| anyhow::anyhow!("Failed to get tools: {}", e))?;
                    if !tools.is_empty() {
                        // Note: In a real implementation, you would need to convert ToolDefinition
                        // to rig's tool format. This is a placeholder for the actual integration.
                        // agent_builder = agent_builder.tools(tools);
                    }
                }
            }

            let agent = agent_builder.build();

            let rig_history: Vec<rig::completion::Message> = history
                .into_iter()
                .map(|msg| {
                    if msg.role == "user" {
                        rig::completion::Message::User {
                            content: OneOrMany::one(UserContent::text(msg.content)),
                        }
                    } else {
                        rig::completion::Message::Assistant {
                            id: None,
                            content: OneOrMany::one(msg.content.into()),
                        }
                    }
                })
                .collect();

            let response = agent
                .chat(prompt, rig_history)
                .await
                .map_err(|e| anyhow::anyhow!("Failed to chat with agent: {}", e))?;

            Ok(response)
        }
    )
}

pub async fn stream_chat(
    agent_manager: &AgentManager,
    tool_manager: &ToolManager,
    agent_id: i64,
    prompt: &str,
    history: Vec<ChatMessage>,
) -> anyhow::Result<Pin<Box<dyn Stream<Item = anyhow::Result<String>> + Send>>> {
    let agent_config = agent_manager
        .get_agent(agent_id)
        .await
        .map_err(|e| anyhow::anyhow!("Failed to get agent: {}", e))?
        .ok_or_else(|| anyhow::anyhow!("Agent with id {} not found", agent_id))?;

    with_client!(
        &agent_config.base_config.provider,
        agent_config.base_config.base_url.as_deref(),
        |client| {
            let agent_builder = client
                .agent(&agent_config.base_config.model)
                .preamble(&agent_config.base_config.preamble);

            // Add tools if available (including MCP tools)
            if let Some(tool_names) = &agent_config.base_config.tools {
                if !tool_names.is_empty() {
                    let tools = tool_manager
                        .get_agent_all_tools(agent_id)
                        .map_err(|e| anyhow::anyhow!("Failed to get tools: {}", e))?;
                    if !tools.is_empty() {
                        // Note: In a real implementation, you would need to convert ToolDefinition
                        // to rig's tool format. This is a placeholder for the actual integration.
                        // agent_builder = agent_builder.tools(tools);
                    }
                }
            }

            let agent = agent_builder.build();

            let rig_history: Vec<rig::completion::Message> = history
                .into_iter()
                .map(|msg| {
                    if msg.role == "user" {
                        rig::completion::Message::User {
                            content: OneOrMany::one(UserContent::text(msg.content)),
                        }
                    } else {
                        rig::completion::Message::Assistant {
                            id: None,
                            content: OneOrMany::one(msg.content.into()),
                        }
                    }
                })
                .collect();

            let full_response = agent
                .chat(prompt, rig_history)
                .await
                .map_err(|e| anyhow::anyhow!("Failed to chat with agent: {}", e));

            let stream =
                futures::stream::once(async move { full_response }).flat_map(
                    |result| match result {
                        Ok(response) => {
                            let chunks: Vec<String> = response
                                .split_whitespace()
                                .map(|word| word.to_string() + " ")
                                .collect();

                            Either::Left(futures::stream::iter(chunks.into_iter().map(Ok)))
                        }
                        Err(e) => Either::Right(futures::stream::once(async move { Err(e) })),
                    },
                );

            Ok(Box::pin(stream))
        }
    )
}
